#pragma once

#include <absl/time/clock.h>
#include <zmq.hpp>

#include "zsocket.hpp"
#include "mdp.h"

const static int g_heartbeat_liveness = 10;

class mdworker
{
public:
	mdworker(std::string broker, std::string service, bool verbose = false) :
		m_worker(m_ctx, zmq::socket_type::dealer),
		m_broker(broker),
		m_service(service),
		m_verbose(verbose),
		m_expect_reply(false),
		m_liveness(g_heartbeat_liveness),
		m_heartbeat(absl::Milliseconds(10000)),
		m_reconnect(absl::Milliseconds(2500))
	{
		connect_to_broker();
	}
public:
	void connect_to_broker()
	{
		m_worker = zsocket(m_ctx, zmq::socket_type::dealer);
		m_worker.connect(m_broker);
		if (m_verbose) {
			info("connecting to broker at {}", m_broker);
		}

		send_to_broker(mdpw_command::ready, m_service);
		m_liveness = g_heartbeat_liveness;
		m_heartbeat_at = absl::Now() + m_heartbeat;
	}
	
	void set_heartbeat(absl::Duration heartbeat) {
		m_heartbeat = heartbeat;
	}

	void set_reconnect(absl::Duration reconnect) {
		m_reconnect = reconnect;
	}

	void send_to_broker(const mdpw_command command, const std::string& option, const zframe& data )
	{
		zframe reply = data.copy();
		if (!option.empty()) {
			reply.push_front(option);
		}
		reply.push_front(int32_t(command));
		reply.push_front(g_mdpw_worker);
		reply.push_front();

		if (m_verbose) {
			info("sending {} to broker", g_mdps_commands[int32_t(command)]);
			reply.dump();
		}
		m_worker.send(reply);
	}

	void send_to_broker(const mdpw_command command, const std::string& option = "")
	{
		zframe reply;
		send_to_broker(command, option, reply);
	}

	zframe receive(const zframe& reply)
	{
		if (!reply.is_empty())
		{
			zframe reply_cp = reply.copy();
			reply_cp.push_front();
			reply_cp.push_front(m_reply_to);
			send_to_broker(mdpw_command::reply, "", std::move(reply_cp));
		}
		m_expect_reply = true;
		while (true)
		{
			zmq::pollitem_t item = { m_worker, 0, ZMQ_POLLIN, 0 };
			zmq::poll(&item, 1, int(absl::ToInt64Milliseconds(m_heartbeat)));
			if (item.revents & ZMQ_POLLIN)
			{
				zframe msg;
				m_worker.recv(msg);
				if (msg.is_empty()) {
					break;
				}

				if (m_verbose) {
					info("received message from broker:");
					msg.dump();
				}
				m_liveness = g_heartbeat_liveness;

				assert(msg.size() >= 3);
				zmsg empty = msg.pop_front();
				zmsg header = msg.pop_front();
				assert(header == g_mdpw_worker);
				zmsg cmd_msg = msg.pop_front();
				mdpw_command cmd = mdpw_command(to_int32(cmd_msg));
				switch (cmd)
				{
				case mdpw_command::request:
					m_reply_to = msg.pop_front();
					msg.pop_front();
					return msg;
					break;
				case mdpw_command::heartbeat:
					break;
				case mdpw_command::disconnect:
					connect_to_broker();
					break;
				default:
					info("invalid input message");
					msg.dump();
					break;
				}
			}
			else if (--m_liveness == 0)
			{
				info("disconnected from broker - retrying...");
				std::this_thread::sleep_for(absl::ToChronoMilliseconds(m_reconnect));
				connect_to_broker();
			}

			if (absl::Now() > m_heartbeat_at)
			{
				send_to_broker(mdpw_command::heartbeat);
				m_heartbeat_at = absl::Now() + m_heartbeat;
			}
		}

		return zframe();
	}


private:
	zmq::context_t m_ctx;
	zsocket m_worker;
	std::string m_broker;
	std::string m_service;
	
private:
	bool m_verbose;
	absl::Time m_heartbeat_at;
	size_t m_liveness;
	absl::Duration m_heartbeat;
	absl::Duration m_reconnect;

private:
	bool m_expect_reply;
	zmsg m_reply_to;
};